# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig
import numpy as np

@GENERATOR_REGISTRY.register("ascend_moeinitroutingv2grad")
class Upsamplebicubic2dGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.tensor_dim = 0
        self.range_is_null = False
        self.input_shape = []
        self.input_dtype = None

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
        '''
        当参数之间有相互依赖关系时，需要覆写此函数来约束生成的输入参数信息
        :param index: 用例参数列表的下标，0：表示第0个参数
        :param input_case: 随机生成的用例信息
        :return: 修改后需要返回的输入信息对象
        '''

        # 限制第一个参数self的shape
        if index == 0:
            #input_case.shape = (8,2,5102)
            self.input_shape = input_case.shape # input shape [N, C]
            self.input_dtype = input_case.dtype

        # 获取第二个参数outputsize的range_values，第一个值表示采样后输出tensor的H值，第二个值表示采样后输出tensor的W值
        if index == 1:
            input_case.shape[0] = self.input_shape[0] # target [N,]; 值域 [0, C) int64
            pass

        return input_case